package com.lyy.mybatis.sensitive.interceptor;


import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.Pair;
import cn.hutool.core.util.ReflectUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.core.util.TypeUtil;
import com.baomidou.mybatisplus.annotation.TableField;
import com.baomidou.mybatisplus.core.conditions.AbstractWrapper;
import com.baomidou.mybatisplus.core.conditions.segments.MergeSegments;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.lyy.mybatis.sensitive.annotation.SensitiveField;
import com.lyy.mybatis.sensitive.mapping.ColumMapping;
import com.lyy.mybatis.sensitive.mapping.FieldMapping;
import com.lyy.mybatis.sensitive.secret.SecretAbstract;
import com.lyy.mybatis.sensitive.utils.EncryptDecryptUtil;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.ExpressionVisitorAdapter;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.LikeExpression;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import org.apache.ibatis.annotations.Param;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.jetbrains.annotations.NotNull;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.context.annotation.Lazy;

import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.lang.reflect.Type;
import java.sql.Connection;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 查询条件加密拦截器
 * 场景：getByPhone => where phone = 密文，入参时会将明文转换成密文以匹配条件
 * 支持查询的方式：
 * 1. mybatis-plus单表查询：xxService.lambdaQuery().xx 、xxService.query().xx
 * 2. 原生mybatis mapper 查询，入参为单一对象的方法： User query(xxDTO dto)  入参对象实现Encrypt
 * 3. 原生mybatis mapper 查询，入参为字符串： User queryByPhone(@Encrypt @Param("phone") String phone)
 *
 * @author lyy 2024/11/23
 */
@Slf4j
@Lazy
@AutoConfiguration
@Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
public class SensitiveQueryInterceptor extends SensitiveBaseInterceptor implements Interceptor {
    /**
     * 缓存：<实体类classPath,Pair<实体类对应需要加密的字段,加密类型>>
     */
    private static final Map<String, List<Pair<String, Class<? extends SecretAbstract>>>> ENCRYPT_FIELD_CACHE = new ConcurrentHashMap<>();
    /**
     * 缓存：<实体类classPath,Pair<MapperClass, EntityClass>>
     */
    private static final Map<String, Pair<Class, Class>> MAPPED_CLASS_CACHE = new ConcurrentHashMap<>();
    /**
     * 缓存：<mapper方法名，Pair<需要加密的参数名,加密类型>> （特指参数是String类型的方法）
     */
    private static final Map<String, List<Pair<String, Class<? extends SecretAbstract>>>> MAPPED_METHOD_CACHE = new ConcurrentHashMap<>();

    @Override
    public void setProperties(Properties properties) {

    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof StatementHandler) {
            return Plugin.wrap(target, this);
        }
        return target;
    }


    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        StatementHandler statementHandler = PluginUtils.realTarget(invocation.getTarget());
        MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
        if (!isSelect(mappedStatement)) {
            return invocation.proceed();
        }
        Object parameterObject = metaObject.getValue("delegate.boundSql.parameterObject");

        if (EncryptDecryptUtil.checkIsSensitiveTable(parameterObject)) {
            //处理原生Mapper+对象入参查询
            this.handleMapperObjectQuery(parameterObject);
        } else if (parameterObject instanceof MapperMethod.ParamMap) {
            MapperMethod.ParamMap paramMap = (MapperMethod.ParamMap) parameterObject;
            if (paramMap.containsKey("ew")) {
                //处理mybatis-plus单表查询
                this.handleMybatisPlusQuery(mappedStatement, parameterObject);
            }
            //处理原生Mapper+字符串入参查询
            else {
                this.handleMapperParamQuery(mappedStatement, parameterObject);
            }
        }
        return invocation.proceed();
    }

    /**
     * 处理原生Mapper+对象入参查询
     */
    @SneakyThrows
    private void handleMapperObjectQuery(Object parameterObject) {
        EncryptDecryptUtil.encrypt(parameterObject);
    }

    /**
     * 处理原生Mapper+字符串入参查询
     */
    @SneakyThrows
    private void handleMapperParamQuery(MappedStatement mappedStatement, Object parameterObject) {
        List<Pair<String, Class<? extends SecretAbstract>>> encryptParamList = getEncryptParamByMethod(mappedStatement);
        MapperMethod.ParamMap<Object> paramMap = (MapperMethod.ParamMap) parameterObject;
        for (Map.Entry<String, Object> param : paramMap.entrySet()) {
            Pair<String, Class<? extends SecretAbstract>> pair = encryptParamList.stream().filter(e -> e.getKey().equals(param.getKey()))
                    .findFirst().orElse(null);
            if (pair != null && param.getValue() != null) {
                paramMap.put(param.getKey(), EncryptDecryptUtil.encryptString(pair.getKey(), pair.getValue(), null));
            }
        }
    }

    /**
     * 处理mybatis-plus单表查询
     */
    private void handleMybatisPlusQuery(MappedStatement mappedStatement, Object parameterObject) {
        List<Pair<String, Class<? extends SecretAbstract>>> encryptFields = getEntityEncryptFields(mappedStatement);
        if (CollUtil.isEmpty(encryptFields)) {
            return;
        }
        MapperMethod.ParamMap paramMap = (MapperMethod.ParamMap) parameterObject;
        AbstractWrapper ew = (AbstractWrapper) paramMap.get("ew");
        MergeSegments expression = ew.getExpression();
        if (expression != null) {
            String whereSql = expression.getSqlSegment();
            Map<String, Object> paramNameValuePairs = ew.getParamNameValuePairs();
            Map<String, ColumMapping> columMappingMap = analysisWhereSql(whereSql, paramNameValuePairs);
            for (Pair<String, Class<? extends SecretAbstract>> encryptField : encryptFields) {
                ColumMapping columMapping = columMappingMap.get(encryptField.getKey());
                if (columMapping != null) {
                    String encryptValue = EncryptDecryptUtil.encryptString(columMapping.getVal(), encryptField.getValue(),
                            columMapping.getExpression());
                    paramNameValuePairs.put(columMapping.getAliasName(), encryptValue);
                }
            }
        }
    }


    /**
     * 解析 whereSql
     *
     * @param whereSql
     * @param paramNameValuePairs
     * @return
     */
    private Map<String, ColumMapping> analysisWhereSql(String whereSql, Map<String, Object> paramNameValuePairs) {
        Map<String, FieldMapping> aliasNameMap = parseParamAliasNameMap(whereSql);
        Map<String, ColumMapping> columMap = new HashMap<>();
        for (Map.Entry<String, Object> item : paramNameValuePairs.entrySet()) {
            //参数别名
            String aliasName = item.getKey();
            //参数值
            String val = item.getValue() != null ? item.getValue().toString() : null;
            //字段映射信息
            FieldMapping fieldMapping = aliasNameMap.get(aliasName);
            columMap.put(
                    fieldMapping.getFieldName(),
                    new ColumMapping()
                            .setAliasName(aliasName)
                            .setVal(val)
                            .setExpression(fieldMapping.getExpression())
            );
        }

        return columMap;
    }

    /**
     * 获取某方法需要加密的入参
     *
     * @return
     */
    private List<Pair<String, Class<? extends SecretAbstract>>> getEncryptParamByMethod(MappedStatement mappedStatement) {
        String id = mappedStatement.getId();
        List<Pair<String, Class<? extends SecretAbstract>>> list = MAPPED_METHOD_CACHE.get(id);
        if (CollUtil.isEmpty(list)) {
            String[] split = id.split("\\.");
            String methodName = split.length != 0 ? split[split.length - 1] : id;
            Class clazz = getEntityClass(mappedStatement).getKey();
            Method method = ReflectUtil.getMethodByName(clazz, methodName);
            list = new ArrayList<>();
            for (Parameter p : method.getParameters()) {
                if (p.isAnnotationPresent(SensitiveField.class) && p.getType() == String.class) {
                    SensitiveField sensitiveField = p.getAnnotation(SensitiveField.class);
                    Class<? extends SecretAbstract> secret = sensitiveField.value();
                    Param param = p.getAnnotation(Param.class);
                    list.add(param != null ? Pair.of(param.value(), secret) : Pair.of(p.getName(), secret));
                }
            }
            MAPPED_METHOD_CACHE.put(id, list);
        }
        return list;
    }

    /**
     * 根据Mybatis-plus的whereSql获得一个参数别名与字段名的映射map
     *
     * @return 别名, 字段名
     */
    @SneakyThrows
    private Map<String, FieldMapping> parseParamAliasNameMap(String whereSql) {
        whereSql = whereSql
                .replaceAll("#", "")
                .replaceAll("\\$", "")
                .replaceAll("\\(", "")
                .replaceAll("\\)", "")
                .replaceAll("\\{", "")
                .replaceAll("}", "")
                .replaceAll("ew.paramNameValuePairs.", "");

        Map<String, FieldMapping> res = new HashMap<>();
        if (StrUtil.isBlank(whereSql)) {
            return res;
        }
        Expression expression = CCJSqlParserUtil.parseCondExpression(whereSql);

        expression.accept(new ExpressionVisitorAdapter() {
            @Override
            public void visit(AndExpression expr) {
                if (expr.getLeftExpression() instanceof AndExpression) {
                    expr.getLeftExpression().accept(this);
                } else if ((expr.getLeftExpression() instanceof EqualsTo)) {
                    FieldMapping fieldMapping = getColumnName(expr.getLeftExpression());
                    if (fieldMapping == null){
                        return;
                    }
                    res.put(fieldMapping.getAliasName(), fieldMapping);
                }
                FieldMapping fieldMapping = getColumnName(expr.getRightExpression());
                if (fieldMapping == null){
                    return;
                }
                res.put(fieldMapping.getAliasName(), fieldMapping);
            }

            @Override
            public void visit(EqualsTo expr) {
                FieldMapping fieldMapping = getColumnName(expr);
                if (fieldMapping == null){
                    return;
                }
                res.put(fieldMapping.getAliasName(), fieldMapping);
            }

            @Override
            public void visit(LikeExpression expr) {
                FieldMapping fieldMapping = getColumnName(expr);
                if (fieldMapping == null){
                    return;
                }
                res.put(fieldMapping.getAliasName(), fieldMapping);
            }
        });
        return res;
    }


    /**
     * 返回where表达式中，字段名与参数名的映射关系
     *
     * @param expr 表达式
     * @return FieldMap<字段名, 别名, 表达式类型> ，参数别名特指：MPGENVAL1、MPGENVAL2 ...
     */
    private FieldMapping getColumnName(Expression expr) {
        if ((expr instanceof EqualsTo)) {
            EqualsTo equalsTo = (EqualsTo) expr;
            return getFieldMap(expr, (Column) equalsTo.getLeftExpression(), equalsTo.getRightExpression());
        } else if (expr instanceof LikeExpression) {
            LikeExpression likeExpression = (LikeExpression) expr;
            return getFieldMap(expr, (Column) likeExpression.getLeftExpression(), likeExpression.getRightExpression());
        }
        return null;
    }

    @NotNull
    private FieldMapping getFieldMap(Expression expr, Column column, Expression rightExpression) {
        String columnName = column.getColumnName();

        String aliasName;
        if (rightExpression instanceof Column) {
            aliasName = ((Column) rightExpression).getColumnName();
        } else {
            aliasName = rightExpression.toString();
        }
        return new FieldMapping(columnName, aliasName, expr);
    }


    /**
     * 获取本次查询的实体类所需要加密的字段名list
     *
     * @param mappedStatement
     * @return
     */
    private List<Pair<String, Class<? extends SecretAbstract>>> getEntityEncryptFields(MappedStatement mappedStatement) {
        String id = mappedStatement.getId();
        List<Pair<String, Class<? extends SecretAbstract>>> encryptFields = ENCRYPT_FIELD_CACHE.get(id);
        if (encryptFields == null) {
            Class entityClass = getEntityClass(mappedStatement).getValue();
            encryptFields = getEntityEncryptFields(entityClass);
            ENCRYPT_FIELD_CACHE.put(id, encryptFields);
        }
        return encryptFields;
    }

    private List<Pair<String, Class<? extends SecretAbstract>>> getEntityEncryptFields(Class clazz) {
        Field[] fields = ReflectUtil.getFields(clazz);
        List<Pair<String, Class<? extends SecretAbstract>>> list = new ArrayList<>();
        for (Field field : fields) {
            if (!isEncryptField(field)) {
                continue;
            }
            SensitiveField sensitiveField = field.getAnnotation(SensitiveField.class);
            Class<? extends SecretAbstract> secret = sensitiveField.value();
            TableField tableField = field.getAnnotation(TableField.class);
            if (tableField != null && StrUtil.isNotBlank(tableField.value())) {
                list.add(Pair.of(tableField.value(), secret));
            }
            //驼峰转下划线
            list.add(Pair.of(StrUtil.toUnderlineCase(field.getName()), secret));
        }
        return list;
    }


    /**
     * 获取EntityClass
     *
     * @param mappedStatement
     * @return mapperClass, entityClass
     */
    @SneakyThrows
    private Pair<Class, Class> getEntityClass(MappedStatement mappedStatement) {
        String id = mappedStatement.getId();
        Pair<Class, Class> pair = MAPPED_CLASS_CACHE.get(id);
        if (pair == null) {
            String classPath = this.getMapperClassPath(mappedStatement);
            Class<?> mapperClass = Class.forName(classPath);
            //从Mapper中获取泛型所指向的实体类 "interface xxxMapper extends BaseMapper<xxxEntity>" 中的xxxEntity
            Type typeArgument = TypeUtil.getTypeArgument(mapperClass, 0);
            Class<?> entityClass = (Class) typeArgument;

            pair = Pair.of(mapperClass, entityClass);
            MAPPED_CLASS_CACHE.put(id, pair);
        }
        return pair;
    }
}